package com.zyx.flinkdemo.dataset.operator;

import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.functions.CoGroupFunction;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.operators.CoGroupOperator;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.List;

/**
 * @author Yaxi.Zhang
 * @since 2021/6/30 14:23
 */
public class CoGroupTemplate {
    public static void main(String[] args) throws Exception {
        List<Tuple2<Integer, String>> commodityList = new ArrayList<Tuple2<Integer, String>>();
        List<Tuple2<String, Integer>> orderList = new ArrayList<Tuple2<String, Integer>>();

        commodityList.add(new Tuple2<>(1, "手机"));
        commodityList.add(new Tuple2<>(2, "电脑"));
        commodityList.add(new Tuple2<>(3, "移动电源"));

        orderList.add(new Tuple2<>("张三", 1));
        orderList.add(new Tuple2<>("赵六", 1));
        orderList.add(new Tuple2<>("李四", 2));
        orderList.add(new Tuple2<>("刘七", 3));
        orderList.add(new Tuple2<>("王五", 2));

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        DataSet<Tuple2<Integer, String>> commodityDataSet = env.fromCollection(commodityList);
        DataSet<Tuple2<String, Integer>> orderDataSet = env.fromCollection(orderList);

        DataSet<Tuple2<String, String>> result =
                commodityDataSet.coGroup(orderDataSet)
                        .where("f0")
                        .equalTo("f1")
                        .with(new CustomCoGroupFunction());

        result.print();


        CoGroupOperator
                .CoGroupOperatorSets<Tuple2<Integer, String>, Tuple2<String, Integer>>
                .CoGroupOperatorSetsPredicate.CoGroupOperatorWithoutFunction coGroupOperatorWithoutFunction = commodityDataSet.coGroup(orderDataSet)
                .where("f0")
                .equalTo("f1");



        env.execute("CoGroup Template");
    }

    public static class CustomCoGroupFunction implements CoGroupFunction<Tuple2<Integer, String>, Tuple2<String, Integer>, Tuple2<String, String>> {
        @Override
        public void coGroup(Iterable<Tuple2<Integer, String>> first, Iterable<Tuple2<String, Integer>> second,
                            Collector<Tuple2<String, String>> out) {
            String names = "";
            for (Tuple2<String, Integer> next : second) {
                names = StringUtils.join(names, " ", next.f0);
            }
            for (Tuple2<Integer, String> next : first) {
                out.collect(new Tuple2<>(next.f1, names));
            }
        }
    }
}
